#include "limbo/arp/arp_request.hpp"
#include "limbo/arp/arp_response.hpp"
#include "limbo/arp/arp_table.hpp"
#include "limbo/arp/state.hpp"
#include "limbo/errc.h"
#include "limbo/ethernet/state.hpp"
#include "limbo/icmp/destination_unreachable.hpp"
#include "limbo/icmp/echo_request.hpp"
#include "limbo/icmp/echo_response.hpp"
#include "limbo/icmp/state.hpp"
#include "limbo/ip/state.hpp"
#include "limbo/udp/connection.hpp"
#include "limbo/udp/receiver.hpp"
#include "limbo/udp/state.hpp"

#include "limbo/stack.hpp"
#include "test-utils.h"

using namespace limbo;

using Eth = ethernet::State<void>;
using Arp = arp::State<Eth>;
using ArpRequest = arp::ArpRequest<Arp>;
using ArpResponse = arp::ArpResponse<Arp>;
using IP = ip::State<Eth>;
using ICMP = icmp::State<IP>;
using UDP = udp::State<IP>;
using EchoReq = icmp::EchoRequest<ICMP>;
using EchoRes = icmp::EchoResponse<ICMP>;
using Unreach = icmp::DestinationUnreacheable<ICMP>;
using Receiver = udp::Receiver<UDP>;
using Connection = udp::Connection<UDP>;
using ArpCache = arp::ArpTable<2>;

using MyStack1 =
    Stack<Layer<Receiver, EchoReq, EchoRes, Unreach>, /* app */
          Layer<UDP, ICMP, ArpRequest, ArpResponse>,  /* transport */
          Layer<IP, Arp>,                             /* network */
          Layer<Eth>                                  /* link */
          >;

using MyStack2 =
    Stack<Layer<Connection, EchoReq, EchoRes, Unreach>, /* app */
          Layer<UDP, ICMP, ArpRequest, ArpResponse>,    /* transport */
          Layer<IP, Arp>,                               /* network */
          Layer<Eth>                                    /* link */
          >;

uint8_t mac_raw_1[] = {0x48, 0xba, 0x4e, 0x51, 0x39, 0xbb};
uint8_t mac_raw_2[] = {0x48, 0xba, 0x4e, 0x51, 0x39, 0xbc};

auto ip_1 = make_address("192.168.100.1");
auto ip_2 = make_address("192.168.100.2");
auto ep_1 = ip::IPv4Endpoint{ip_1, 2000};
auto ep_2 = ip::IPv4Endpoint{ip_2, 43210};
auto ep_x = ip::IPv4Endpoint{ip_1, 2001};
auto mac_1 = ethernet::MacAddress(mac_raw_1);
auto mac_2 = ethernet::MacAddress(mac_raw_2);

TEST_CASE("udp connection over ethernet", "[scenario]") {
  char buff_raw[1500];
  auto buff = Chunk(buff_raw, sizeof(buff_raw));
  char ping_raw[] = "ping";
  auto ping = Chunk(ping_raw, sizeof(ping_raw));
  char pong_raw[] = "pong";
  auto pong = Chunk(pong_raw, sizeof(pong_raw));

  MyStack1 stack1;
  MyStack2 stack2;

  auto &eth_2 = stack2.get<3, 0>();
  auto &arp_req_2 = stack2.get<1, 2>();
  auto &arp_res_2 = stack2.get<1, 3>();
  auto &arp_req_1 = stack1.get<1, 2>();
  auto &arp_res_1 = stack1.get<1, 3>();
  auto &udp_recv_1 = stack1.get<0, 0>();
  auto &udp_conn_2 = stack2.get<0, 0>();
  udp_recv_1.init(ep_1);
  udp_conn_2.init(ep_2, ep_1);
  arp_req_1.init(mac_1, ip_1);
  arp_res_1.init(mac_1, ip_1);
  arp_req_2.init(mac_2, ip_2);
  arp_res_2.init(mac_2, ip_2);

  // step 0, attempt to send ping, while arp is not known
  auto r2 = udp_conn_2.send(stack2, buff, ping);
  REQUIRE(!r2);
  REQUIRE(r2.state() == &eth_2);
  CHECK(r2.error_code() == (uint32_t)Errc::ethernet_unknown_mac);

  // step 1, exchange macs
  ArpCache cache1, cache2;

  r2 = arp_req_2.send(stack2, buff, ip_1);
  REQUIRE(r2);

  auto r1 = stack1.recv(r2.consumed(), nullptr);
  REQUIRE(r1);
  REQUIRE(r1.state() == &arp_req_1);
  auto &arp_packet_1 = arp_req_1.get_parsed();
  cache1.update(arp_packet_1.sender_mac, arp_packet_1.sender_ip, 5);

  r1 = arp_res_1.send(stack1, buff, arp_req_1);
  REQUIRE(r1);

  r2 = stack2.recv(r1.consumed(), nullptr);
  REQUIRE(r2);
  REQUIRE(r2.state() == &arp_res_2);
  auto &arp_packet_2 = arp_res_2.get_parsed();
  cache2.update(arp_packet_2.sender_mac, arp_packet_2.sender_ip, 5);

  REQUIRE(*cache1.get(mac_2, 2) == ip_2);
  REQUIRE(*cache2.get(mac_1, 2) == ip_1);

  // step 2, send "pong", receive "ping"
  auto &link_ctx_2 = get_link_context(udp_conn_2);
  link_ctx_2.source = mac_2;
  link_ctx_2.destination = mac_1;
  link_ctx_2.type = ethernet::EtherType::ipv4;
  r2 = udp_conn_2.send(stack2, buff, ping);
  REQUIRE(r2);

  r1 = stack1.recv(r2.consumed(), nullptr);
  REQUIRE(r1);
  REQUIRE(r1.state() == &udp_recv_1);
  auto &udp_packet_1 = udp_recv_1.get_parsed();
  CHECK(udp_packet_1.payload == ping);
  auto &udp_1_mac = get_link_context(udp_packet_1).source;
  auto &udp_1_ip = get_ip_context(udp_packet_1).source;
  cache1.update(udp_1_mac, udp_1_ip, 6);

  auto link_ctx_1 =
      Eth::Context{mac_1, udp_1_mac, ethernet::EtherType::ipv4, nullptr};
  get_ip_context(udp_recv_1).link_context = &link_ctx_1;
  r1 = udp_recv_1.send(stack1, buff, ep_2, pong);
  REQUIRE(r1);

  r2 = stack2.recv(r1.consumed(), nullptr);
  REQUIRE(r2);
  REQUIRE(r2.state() == &udp_conn_2);
  auto &udp_packet_2 = udp_conn_2.get_parsed();
  CHECK(udp_packet_2.payload == pong);

  auto &udp_2_mac = get_link_context(udp_packet_2).source;
  auto &udp_2_ip = get_ip_context(udp_packet_2).source;
  cache2.update(udp_2_mac, udp_2_ip, 6);
}

TEST_CASE("wrong connection to dup, reply back with arp", "[scenario]") {
  char buff_raw[1500];
  auto buff = Chunk(buff_raw, sizeof(buff_raw));
  char buff_raw_2[1500];
  auto buff_2 = Chunk(buff_raw_2, sizeof(buff_raw_2));
  char ping_raw[] = "ping";
  auto ping = Chunk(ping_raw, sizeof(ping_raw));

  MyStack1 stack1;
  MyStack2 stack2;

  auto &arp_req_2 = stack2.get<1, 2>();
  auto &arp_res_2 = stack2.get<1, 3>();
  auto &arp_req_1 = stack1.get<1, 2>();
  auto &arp_res_1 = stack1.get<1, 3>();
  auto &udp_1 = stack1.get<1, 0>();
  auto &udp_conn_2 = stack2.get<0, 0>();
  auto &unreach_1 = stack1.get<0, 3>();
  auto &unreach_2 = stack2.get<0, 3>();
  udp_conn_2.init(ep_2, ep_x);
  arp_req_1.init(mac_1, ip_1);
  arp_res_1.init(mac_1, ip_1);
  arp_req_2.init(mac_2, ip_2);
  arp_res_2.init(mac_2, ip_2);
  unreach_1.init(ip_1);
  unreach_2.init(ip_2);

  // step 1, exchange macs
  ArpCache cache1, cache2;

  auto r2 = arp_req_2.send(stack2, buff, ip_1);
  REQUIRE(r2);

  auto r1 = stack1.recv(r2.consumed(), nullptr);
  REQUIRE(r1);
  REQUIRE(r1.state() == &arp_req_1);
  auto &arp_packet_1 = arp_req_1.get_parsed();
  cache1.update(arp_packet_1.sender_mac, arp_packet_1.sender_ip, 5);

  r1 = arp_res_1.send(stack1, buff, arp_req_1);
  REQUIRE(r1);

  r2 = stack2.recv(r1.consumed(), nullptr);
  REQUIRE(r2);
  REQUIRE(r2.state() == &arp_res_2);
  auto &arp_packet_2 = arp_res_2.get_parsed();
  cache2.update(arp_packet_2.sender_mac, arp_packet_2.sender_ip, 5);

  REQUIRE(*cache1.get(mac_2, 2) == ip_2);
  REQUIRE(*cache2.get(mac_1, 2) == ip_1);

  // step 2, send "pong", receive bare udp
  auto &link_ctx_2 = get_link_context(udp_conn_2);
  link_ctx_2.source = mac_2;
  link_ctx_2.destination = mac_1;
  link_ctx_2.type = ethernet::EtherType::ipv4;
  r2 = udp_conn_2.send(stack2, buff, ping);
  REQUIRE(r2);

  r1 = stack1.recv(r2.consumed(), nullptr);
  REQUIRE(r1);
  REQUIRE(r1.state() == &udp_1);
  auto &udp_packet_1 = udp_1.get_parsed();
  CHECK(udp_packet_1.payload == ping);
  auto &udp_1_mac = get_link_context(udp_packet_1).source;
  auto &udp_1_ip = get_ip_context(udp_packet_1).source;
  cache1.update(udp_1_mac, udp_1_ip, 6);

  // step 3, send & receive ICMP destination unreacheable
  auto &link_ctx_1 = get_link_context(unreach_1);
  link_ctx_1.source = mac_1;
  link_ctx_1.destination = mac_2;
  link_ctx_1.type = ethernet::EtherType::ipv4;
  /* buff_1 cannot be used, as we take ip from here */
  r1 = unreach_1.send(stack1, buff_2, *udp_packet_1.container, 3);
  REQUIRE(r1);

  r2 = stack2.recv(r1.consumed(), nullptr);
  REQUIRE(r2);
  REQUIRE(r2.state() == &unreach_2);
  auto &icmp_packet = unreach_2.get_parsed();
  REQUIRE(icmp_packet.type == icmp::Type::destination_unreachable);
  REQUIRE(icmp_packet.code == 3);
}

TEST_CASE("icmp echo req/res", "[scenario]") {
  char buff_raw[1500];
  auto buff = Chunk(buff_raw, sizeof(buff_raw));
  char ping_raw[] = "ping";
  auto ping = Chunk(ping_raw, sizeof(ping_raw));

  MyStack1 stack1;
  MyStack2 stack2;

  auto &arp_req_2 = stack2.get<1, 2>();
  auto &arp_res_2 = stack2.get<1, 3>();
  auto &arp_req_1 = stack1.get<1, 2>();
  auto &arp_res_1 = stack1.get<1, 3>();
  auto &echo_res_1 = stack1.get<0, 2>();
  auto &echo_res_2 = stack2.get<0, 2>();
  auto &echo_req_1 = stack1.get<0, 1>();
  auto &echo_req_2 = stack2.get<0, 1>();
  arp_req_1.init(mac_1, ip_1);
  arp_res_1.init(mac_1, ip_1);
  arp_req_2.init(mac_2, ip_2);
  arp_res_2.init(mac_2, ip_2);
  echo_req_1.init(ip_1);
  echo_req_2.init(ip_2);
  echo_res_1.init(ip_1);
  echo_res_2.init(ip_2);

  // step 1, exchange macs
  ArpCache cache1, cache2;

  auto r2 = arp_req_2.send(stack2, buff, ip_1);
  REQUIRE(r2);

  auto r1 = stack1.recv(r2.consumed(), nullptr);
  REQUIRE(r1);
  REQUIRE(r1.state() == &arp_req_1);
  auto &arp_packet_1 = arp_req_1.get_parsed();
  cache1.update(arp_packet_1.sender_mac, arp_packet_1.sender_ip, 5);

  r1 = arp_res_1.send(stack1, buff, arp_req_1);
  REQUIRE(r1);

  r2 = stack2.recv(r1.consumed(), nullptr);
  REQUIRE(r2);
  REQUIRE(r2.state() == &arp_res_2);
  auto &arp_packet_2 = arp_res_2.get_parsed();
  cache2.update(arp_packet_2.sender_mac, arp_packet_2.sender_ip, 5);

  REQUIRE(*cache1.get(mac_2, 2) == ip_2);
  REQUIRE(*cache2.get(mac_1, 2) == ip_1);

  // step 2, send & receive echo request
  auto &link_ctx_2 = get_link_context(echo_req_2);
  link_ctx_2.source = mac_2;
  link_ctx_2.destination = mac_1;
  link_ctx_2.type = ethernet::EtherType::ipv4;

  r2 = echo_req_2.send(stack2, buff, ip_1, ping);
  REQUIRE(r2);

  r1 = stack1.recv(r2.consumed(), nullptr);
  REQUIRE(r1);
  REQUIRE(r1.state() == &echo_req_1);
  auto &echo_packet_1 = echo_req_1.get_parsed();
  auto &ip_packet = *echo_packet_1.container->container;
  CHECK(ip_packet.source == ip_2);
  CHECK(ip_packet.destination == ip_1);
  REQUIRE(echo_packet_1.payload == ping);

  // step 3, send & receve echo response
  auto &link_ctx_1 = get_link_context(echo_res_1);
  link_ctx_1.source = mac_1;
  link_ctx_1.destination = mac_2;
  link_ctx_1.type = ethernet::EtherType::ipv4;
  r1 = echo_res_1.send(stack1, buff, echo_packet_1);
  REQUIRE(r1);

  r2 = stack2.recv(r1.consumed(), nullptr);
  REQUIRE(r2);
  REQUIRE(r2.state() == &echo_res_2);

  auto &echo_packet_2 = echo_res_2.get_parsed();
  ip_packet = *echo_packet_2.container->container;
  CHECK(ip_packet.source == ip_1);
  CHECK(ip_packet.destination == ip_2);
  REQUIRE(echo_packet_2.payload == ping);
}
